package com.micro.ai.models.service.impl;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import com.micro.ai.commons.exception.BusinessException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;

/**
 * 数据集处理器
 * 负责解析不同格式的数据集文件并计算质量评分
 * 
 * @author micro-ai
 * @since 0.0.1
 */
@Slf4j
@Component
public class FineTuningDatasetProcessor {

    /**
     * 解析数据集文件
     * 
     * @param filePath 文件路径
     * @param fileFormat 文件格式（jsonl、csv、txt）
     * @return 数据记录列表
     */
    public List<DatasetRecordParseResult> parseDatasetFile(String filePath, String fileFormat) {
        List<DatasetRecordParseResult> records = new ArrayList<>();
        
        try {
            switch (fileFormat.toLowerCase()) {
                case "jsonl":
                    records = parseJSONL(filePath);
                    break;
                case "csv":
                    records = parseCSV(filePath);
                    break;
                case "txt":
                    records = parseTXT(filePath);
                    break;
                case "parquet":
                    records = parseParquet(filePath);
                    break;
                default:
                    throw new BusinessException("M0000", "不支持的文件格式: " + fileFormat);
            }
        } catch (IOException e) {
            log.error("解析数据集文件失败: filePath={}, error={}", filePath, e.getMessage(), e);
            throw new BusinessException("M0003", "解析文件失败: " + e.getMessage());
        }
        
        return records;
    }

    /**
     * 解析JSONL格式（每行一个JSON对象）
     * 支持格式：
     * - {"input": "...", "output": "..."}
     * - {"prompt": "...", "completion": "..."}
     * - {"messages": [...], "response": "..."}
     */
    private List<DatasetRecordParseResult> parseJSONL(String filePath) throws IOException {
        List<DatasetRecordParseResult> records = new ArrayList<>();
        
        try (BufferedReader reader = new BufferedReader(new FileReader(filePath))) {
            String line;
            int index = 0;
            
            while ((line = reader.readLine()) != null) {
                if (line.trim().isEmpty()) {
                    continue;
                }
                
                try {
                    JSONObject json = JSON.parseObject(line);
                    
                    DatasetRecordParseResult record = new DatasetRecordParseResult();
                    record.setRecordIndex(index++);
                    
                    // 支持多种字段名称
                    record.setInputText(
                        json.getString("input") != null ? json.getString("input") :
                        json.getString("prompt") != null ? json.getString("prompt") :
                        json.getString("messages") != null ? json.getJSONArray("messages").toJSONString() :
                        ""
                    );
                    
                    record.setOutputText(
                        json.getString("output") != null ? json.getString("output") :
                        json.getString("completion") != null ? json.getString("completion") :
                        json.getString("response") != null ? json.getString("response") :
                        ""
                    );
                    
                    record.setInstruction(json.getString("instruction"));
                    
                    // 计算单条记录的质量评分
                    record.setQualityScore(calculateRecordQuality(record));
                    record.setIsValid(record.getQualityScore() >= 0.5);
                    
                    records.add(record);
                    
                } catch (Exception e) {
                    log.warn("解析JSONL行失败: line={}, error={}", line, e.getMessage());
                    // 跳过无效行，继续处理
                }
            }
        }
        
        return records;
    }

    /**
     * 解析CSV格式
     */
    private List<DatasetRecordParseResult> parseCSV(String filePath) throws IOException {
        List<DatasetRecordParseResult> records = new ArrayList<>();
        
        List<String> lines = Files.readAllLines(Paths.get(filePath));
        if (lines.isEmpty()) {
            return records;
        }
        
        // 解析表头
        String[] headers = lines.get(0).split(",");
        int inputIndex = findColumnIndex(headers, new String[]{"input", "prompt", "question"});
        int outputIndex = findColumnIndex(headers, new String[]{"output", "completion", "answer", "response"});
        int instructionIndex = findColumnIndex(headers, new String[]{"instruction", "system"});
        
        // 解析数据行
        for (int i = 1; i < lines.size(); i++) {
            String[] values = lines.get(i).split(",");
            
            if (values.length <= Math.max(inputIndex, outputIndex)) {
                continue; // 跳过不完整的行
            }
            
            DatasetRecordParseResult record = new DatasetRecordParseResult();
            record.setRecordIndex(i - 1);
            record.setInputText(inputIndex >= 0 ? values[inputIndex] : "");
            record.setOutputText(outputIndex >= 0 ? values[outputIndex] : "");
            record.setInstruction(instructionIndex >= 0 ? values[instructionIndex] : null);
            record.setQualityScore(calculateRecordQuality(record));
            record.setIsValid(record.getQualityScore() >= 0.5);
            
            records.add(record);
        }
        
        return records;
    }

    /**
     * 解析TXT格式（每两行为一组：输入和输出）
     */
    private List<DatasetRecordParseResult> parseTXT(String filePath) throws IOException {
        List<DatasetRecordParseResult> records = new ArrayList<>();
        
        List<String> lines = Files.readAllLines(Paths.get(filePath));
        
        for (int i = 0; i < lines.size(); i += 2) {
            DatasetRecordParseResult record = new DatasetRecordParseResult();
            record.setRecordIndex(records.size());
            
            if (i < lines.size()) {
                record.setInputText(lines.get(i));
            }
            if (i + 1 < lines.size()) {
                record.setOutputText(lines.get(i + 1));
            }
            
            record.setQualityScore(calculateRecordQuality(record));
            record.setIsValid(record.getQualityScore() >= 0.5);
            
            records.add(record);
        }
        
        return records;
    }

    /**
     * 解析Parquet格式（需要额外依赖）
     */
    private List<DatasetRecordParseResult> parseParquet(String filePath) {
        // TODO: 实现Parquet文件解析
        // 需要使用org.apache.parquet或类似库
        throw new BusinessException("M0000", "Parquet格式暂未实现，请使用JSONL或CSV格式");
    }

    /**
     * 计算单条记录的质量评分
     */
    private double calculateRecordQuality(DatasetRecordParseResult record) {
        double score = 1.0;
        
        // 检查输入是否为空
        if (!StringUtils.hasText(record.getInputText())) {
            score -= 0.5;
        }
        
        // 检查输出是否为空
        if (!StringUtils.hasText(record.getOutputText())) {
            score -= 0.5;
        }
        
        // 检查输入长度是否合理（过长可能有问题）
        if (record.getInputText() != null) {
            int inputLength = record.getInputText().length();
            if (inputLength == 0) {
                score -= 0.5;
            } else if (inputLength > 50000) {
                score -= 0.2; // 过长
            } else if (inputLength < 10) {
                score -= 0.1; // 过短
            }
        }
        
        // 检查输出长度是否合理
        if (record.getOutputText() != null) {
            int outputLength = record.getOutputText().length();
            if (outputLength == 0) {
                score -= 0.5;
            } else if (outputLength > 100000) {
                score -= 0.2; // 过长
            }
        }
        
        // 检查是否包含特殊字符（可能是编码问题）
        if (record.getInputText() != null && containsInvalidChars(record.getInputText())) {
            score -= 0.1;
        }
        
        if (record.getOutputText() != null && containsInvalidChars(record.getOutputText())) {
            score -= 0.1;
        }
        
        return Math.max(0.0, Math.min(1.0, score));
    }

    /**
     * 检查是否包含无效字符
     */
    private boolean containsInvalidChars(String text) {
        // 检查是否包含过多的控制字符
        long controlCharCount = text.chars()
            .filter(c -> Character.isISOControl(c) && c != '\n' && c != '\r' && c != '\t')
            .count();
        
        return controlCharCount > text.length() * 0.01; // 超过1%的控制字符
    }

    /**
     * 查找列索引
     */
    private int findColumnIndex(String[] headers, String[] possibleNames) {
        for (String name : possibleNames) {
            for (int i = 0; i < headers.length; i++) {
                if (headers[i].trim().equalsIgnoreCase(name)) {
                    return i;
                }
            }
        }
        return -1;
    }

    /**
     * 计算整体数据集的质量评分
     */
    public double calculateDatasetQualityScore(List<DatasetRecordParseResult> records) {
        if (records.isEmpty()) {
            return 0.0;
        }
        
        double totalScore = 0.0;
        int validCount = 0;
        
        for (DatasetRecordParseResult record : records) {
            if (record.getIsValid() != null && record.getIsValid()) {
                totalScore += record.getQualityScore();
                validCount++;
            }
        }
        
        if (validCount == 0) {
            return 0.0;
        }
        
        // 计算平均值，并考虑有效记录比例
        double avgScore = totalScore / validCount;
        double validRatio = (double) validCount / records.size();
        
        return avgScore * validRatio;
    }

    /**
     * 数据集记录解析结果
     */
    public static class DatasetRecordParseResult {
        private int recordIndex;
        private String inputText;
        private String outputText;
        private String instruction;
        private Double qualityScore;
        private Boolean isValid;

        // Getters and Setters
        public int getRecordIndex() {
            return recordIndex;
        }

        public void setRecordIndex(int recordIndex) {
            this.recordIndex = recordIndex;
        }

        public String getInputText() {
            return inputText;
        }

        public void setInputText(String inputText) {
            this.inputText = inputText;
        }

        public String getOutputText() {
            return outputText;
        }

        public void setOutputText(String outputText) {
            this.outputText = outputText;
        }

        public String getInstruction() {
            return instruction;
        }

        public void setInstruction(String instruction) {
            this.instruction = instruction;
        }

        public Double getQualityScore() {
            return qualityScore;
        }

        public void setQualityScore(Double qualityScore) {
            this.qualityScore = qualityScore;
        }

        public Boolean getIsValid() {
            return isValid;
        }

        public void setIsValid(Boolean isValid) {
            this.isValid = isValid;
        }
    }
}

